import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf as pacf


# 因子分布箱型图
def plot_boxplot(factors):
    """因子分布箱型图"""
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.boxplot([factors[col] for col in factors],
               showmeans=False,
               labels=factors.columns)
    plt.show()


# 因子自相关性
def plot_acf(factor):
    """因子自相关性"""
    pacf(factor)
    plt.show()


# 计算IC
def get_ic(factors, prices, n=1, ret_all=True):
    """计算IC"""
    ret = prices.shift(-n) / prices - 1
    corr = pd.concat([ret, factors], axis=1).corr()
    if not ret_all:
        return corr.iloc[:, 0].to_frame().iloc[1:]
    return corr


# 计算rolling IC
def get_roll_ic(factor, prices, n=1, win=50):
    """计算rolling IC"""
    ret = prices.shift(-n) / prices - 1
    res = factor.rolling(win).agg(lambda x: pd.concat([ret[x.index], x], axis=1).corr().iloc[1,0])
    return res


# IC分布图
def plot_hist(ic, num_bins=50, xtitle='ic'):
    """IC分布图"""
    num_bins = num_bins

    fig, ax = plt.subplots()

    # the histogram of the data
    n, bins, patches = ax.hist(ic, num_bins, density=True)
    mu = ic.mean()  # mean of distribution
    sigma = ic.std()  # standard deviation of distribution

    # add a 'best fit' line
    y = ((1 / (np.sqrt(2 * np.pi) * sigma)) *
         np.exp(-0.5 * (1 / sigma * (bins - mu))**2))
    ax.plot(bins, y, '--')
    ax.set_xlabel(xtitle)
    ax.set_ylabel('Probability density')
    ax.set_title(rf'Histogram of IC: $\mu={mu}$, $\sigma={sigma}$')
    fig.tight_layout()
    plt.show()


# IC 折线图
def plot_line(factor):
    """折线图"""
    fig, ax = plt.subplots(figsize=(12, 6))
    factor.plot(ax=ax)
    plt.show()


# IC热力图
def plot_heatmap(ics):
    """热力图"""
    f, ax = plt.subplots(figsize=(20, 18))
    sns.heatmap(ics, annot=True,  linewidths=.5, ax=ax)
    plt.show()
